#!/usr/bin/env python2
#coding=utf-8

import sys, urllib, urllib2, cookielib
from PyQt4.QtCore import *
from PyQt4.QtWebKit import *
import re, os
from pyquery.pyquery import PyQuery

reload(sys)
sys.setdefaultencoding('utf-8')

'''
使用QWebkit模拟js执行
xssresult 弹窗的结果
data url或html
baseurl 如果data为html则应设置基本url
'''	
class Render(QWebPage):
    def __init__(self):
        # self.app = QApplication(sys.argv)
        QWebPage.__init__(self)
        self.result = False
        self.XSSResult = False
        self.loadFinished.connect(self._loadFinished)

    def browser(self, XSSResult, data, url=True, baseUrl=False):
        self.result = False
        self.XSSResult = XSSResult

        if url:
            self.mainFrame().load(QUrl(data))
            # return self.app.exec_()
        else:
            if baseUrl:
                self.mainFrame().setHtml(data, QUrl(baseUrl))
            else:
                self.mainFrame().setHtml(data)

    def _loadFinished(self, result):
        self.frame = self.mainFrame()
        # self.app.quit()

    def getHtml(self):
        return self.frame.toHtml().toUtf8()

    def javaScriptAlert(self, frame, message):
        if not self.result and message == self.XSSResult:
            self.result = True
        else:
            self.result = False

''' 
功能: 生成检查的数据
参数:
    url: 操作的url
    options: [url_parameters, method]
''' 
class ParseData():
    def __init__(self, url, **options):
        self.url = url
        self.options = options
        self.parameters = None
        self.form = False
        # self.error = False

        opener = urllib2.build_opener()
        try:
            self.html = opener.open(url).read()
        except urllib2.HTTPError as e:
            print 'ParseData error: con\'t open url %s, statusCode:%d'%(url, e.code)
            self.html = '<html><head></head><body></body></html>'
        except Exception as e:
            self.html = '<html><head></head><body></body></html>'

        self._urlParams()
        self._formParam()
    
    def _urlParams(self):
        if 'url_parameter' in self.options:
            return self.options.get('url_parameter')

        if self.url.find('?') > 0:
            self.parameters = []
            paramString = self.url[self.url.find('?')+1:]
            paramTuple = paramString.split('&')
            for p in paramTuple:
                self.parameters.append(p.split('=')[0])
        else:
            self.parameters = False


    def _formParam(self):
        
        if ('method' in self.options and self.options.get('method') == 'post') or self.hasForm():
            pq = PyQuery(self.html)
            forms = pq('form')    
            self.form = []
            
            for i in range(forms.length):
                inputs = []
                formTag = forms.eq(i)
                pinputs = formTag('input')
                action = formTag.attr('action')
                
                url = self.url
                frames = formTag.parents('iframe')+formTag.parents('frame')
                base = pq('base')
                if base:
                    url = base.attr('href')
                elif frames:
                    url = frames[0].attr('src')
                absulateUrlPattern = re.compile(
                    r"(((:?http|https)://)?"
                    r"(([a-zA-Z0-9\._-]+\.[a-zA-Z]{2,6})|([0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3})))"
                    , re.IGNORECASE)

                if action and re.match(r'(https?://)', action) == None:
                    if action and action[0] == '/':
                        action = absulateUrlPattern.search(url).group(0)+action
                    else:
                        if url[len(url)-1] == '/' or re.search(r"(\.\S[^/]+)$", url) != None:
                            # print re.search(r'(\S+)/', url).group(1)
                            action = re.search(r'(\S+)/', url).group(0) + formTag.attr['action']
                        else:
                            action = url + '/' + formTag.attr['action']
                
                name = False
                for j in range(pinputs.length): 
                    if name == pinputs.eq(j).attr('name') or pinputs.eq(j).attr('name') == None:
                        continue

                    inputs.append({
                        'type': pinputs.eq(j).attr('type'),
                        'name': pinputs.eq(j).attr('name'),
                        'value': pinputs.eq(j).attr('value'),
                    })
                    name = pinputs.eq(j).attr('name')

                form = {
                    'action': action,
                    'inputs': inputs,
                    'method': formTag.attr('method'),
                }
                self.form.append(form)

    def hasForm(self):
        pq = PyQuery(self.html)
        if pq.find('form'):
            return True
        return False

    def get(self, payload):
        checkUrl = []

        def matchCallback(match):
            repl = match.group(1)+match.group(2)+payload+match.group(4)
            return repl

        for param in self.parameters:
            patern = re.compile('(\?|&)('+param+'=)(\S+?)(&|$)')
            checkUrl.append(re.sub(patern, matchCallback, self.url, 1))

        return checkUrl

    def post(self, payload):
        checkData = []
        for form in self.form:
            data = []
            p = {}
            for inputs in form['inputs']:
                if inputs['value']:
                    p[inputs['name']] = inputs['value']
                else:
                    p[inputs['name']] = payload

            data.append(p)
            p = {}

            checkData.append({
                'action': form['action'],
                'method': form['method'],
                'data': data
            })
        return checkData

class CheckXSS():
    def __init__(self, url, xss):
        self.r = Render()
        self.p = ParseData(url)
        self.url = url
        self.result = []

        # match = re.search('^((?:http|ftp)s?://([\S^/]+?))/', url)
        # if match:
        #     self.baseUrl = match.group(0)
        # else:
        self.baseUrl = False

        self.cookie = False
        self.getCookie()

        if type(xss) is dict: 
            self.checkGet(xss)
            self.checkPost(xss)

        elif type(xss) is list:
            for payload in xss:
                self.checkGet(payload)
                self.checkPost(payload)
        else:
            print '参数类型错误!'

    def checkGet(self, xss):
        #print self.p.parameters
        if not self.p.parameters:
            return
        checkUrl = self.p.get(xss['code'])
        for u in checkUrl:
            self.r.browser(xss['result'], u)
            print "GET %d"%self.r.result
            if self.r.result:
                    self.generalResult()

    def checkPost(self, xss):
        if not self.p.form:
            return

        checkData = self.p.post(xss['code'])
        for d in checkData:
            for data in d['data']:
                if not d['action']:
                    action = self.url
                else:
                    action = d['action']

                if d['method'] and d['method'].lower() == 'post':
                    statuscode, f = self.open(action, data)
                else:
                    # print action+'?'+urllib.urlencode(data)
                    action = action+'?'+urllib.urlencode(data)
                    statuscode, f = self.open(action)

                self.r.browser(xss['result'], f, False, self.baseUrl)
                print "POST %d"%self.r.result
                if self.r.result:
                    if statuscode == 200:
                        self.generalResult('POST', xss['code'], self.r.result)
                    else:
                        self.generalResult('POST', xss['code'], 2)

    def getCookie(self):
        cookiePath = os.path.dirname(__file__)+'/../runtime/cookie.txt'
        cookie = cookielib.MozillaCookieJar(cookiePath)
        handler = urllib2.HTTPCookieProcessor(cookie)
        opener = urllib2.build_opener(handler)
        try:
            response = opener.open(self.url)
        # except urllib2.HTTPError as e:
        #     print e
        except Exception as e:
            print e
        cookie.save(ignore_discard=True, ignore_expires=True)
        self.cookie = cookie

    def open(self, url, data=False):
        opener = urllib2.build_opener(urllib2.HTTPCookieProcessor(self.cookie))

        result = False
        try:
            if data:
                postdata = urllib.urlencode(data)
                result = opener.open(url, postdata)
            else:
                result = opener.open(url)
        except urllib2.HTTPError as e:
            print e
            return e.code, ''
        except Exception as e:
            print e

        self.cookie.save(ignore_discard=True, ignore_expires=True)

        if result != False:
            return result.getcode(), result.read()

    def generalResult(self, method, xss, result):
        res = {
            'method': method,
            'xss': xss
        }

        if result == 0:
            res['result'] = 'safe'
        elif result == 1:
            res['result'] = 'danger'
        else:
            res['result'] = 'unknown'
        self.result.append(res)
